#!/usr/bin/env python3

"""
LeDog ZMQ 主机 ROS2 节点。
在 OPI 上运行，作为 ROS 节点：
1.  (ZMQ PULL): 接收来自 PC (LeRobot) 的动作指令 (语义键名)。
2.  (ROS Pub):   将指令翻译并发布到:
    - /alphadog_node/set_velocity (SetVelocity)
    - /alphadog_node/set_rpy (SetRpy)
    - /arm_controller/commands (Float64MultiArray)
    - /gripper_controller/commands (Float64MultiArray)
3.  (ROS Sub):   订阅 OPI 上的状态/图像话题。
4.  (ROS Sub):   订阅 /alphadog_node/body_status (BodyStatus) 获取机器狗状态。
5.  (ROS Sub):   将 /joint_states (数字键名) 翻译回语义键名。
6.  (ZMQ PUSH):  将观测数据打包并通过 ZMQ 发送回 PC。
"""

import base64
import json
import logging
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Dict

import cv2
import numpy as np
import zmq

# ROS 依赖
try:
    import rclpy
    from cv_bridge import CvBridge
    from rclpy.node import Node
    from sensor_msgs.msg import Image, JointState
    from std_msgs.msg import Float64MultiArray
    from ros_alphadog.msg import SetVelocity, SetRpy, BodyStatus
except ImportError:
    print("ROS2 (rclpy, sensor_msgs, etc.) 或 cv_bridge 未安装。")
    print("请确保您在 OPI 的 ROS 环境中运行此脚本，并已 sourcing 您的 alphadog 工作区。")
    exit(1)


logging.basicConfig(level=logging.INFO)


@dataclass
class LeDogHostConfig:
    """OPI 主机脚本的配置。"""
    port_zmq_cmd: int = 5555
    port_zmq_observations: int = 5556
    connection_time_s: int = 3600 * 24
    watchdog_timeout_ms: int = 500
    max_loop_freq_hz: int = 30
    camera_names: list = field(default_factory=lambda: ["front", "wrist"])


class LeDogRosBridge(Node):
    """
    一个 ROS2 节点，用于桥接 ZMQ 命令和 ROS 硬件接口。
    """
    def __init__(self, camera_names: list):
        super().__init__("ledog_zmq_host_node")
        self.get_logger().info("LeDog ZMQ-ROS 桥接节点正在启动...")
        self.bridge = CvBridge()

        # ----- 状态变量 (线程安全) -----
        self.lock = threading.Lock()
        
        # --- 翻译映射 ---
        self.pc_arm_joint_keys = [
            "arm_shoulder_pan.pos", "arm_shoulder_lift.pos", "arm_elbow_flex.pos",
            "arm_wrist_flex.pos", "arm_wrist_roll.pos"
        ]
        self.pc_gripper_key = "arm_gripper.pos"
        self.ros_arm_joint_names = ["1", "2", "3", "4", "5"]
        self.ros_gripper_joint_name = "6"
        self.ros_to_pc_map = dict(zip(self.ros_arm_joint_names, self.pc_arm_joint_keys))
        self.ros_to_pc_map[self.ros_gripper_joint_name] = self.pc_gripper_key
        self.dog_control_keys = [
            "dog.vx", "dog.vy", "dog.wz", 
            "dog.pitch", "dog.roll", "dog.yaw"
        ]
        
        self._latest_observation = {
            **{key: 0.0 for key in self.pc_arm_joint_keys},
            self.pc_gripper_key: 0.0,
            **{key: 0.0 for key in self.dog_control_keys}
        }
        self._latest_frames: Dict[str, np.ndarray] = {}

        # ----- ROS 动作发布者 -----
        self.dog_velocity_pub = self.create_publisher(SetVelocity, '/alphadog_node/set_velocity', 10)
        self.dog_rpy_pub = self.create_publisher(SetRpy, '/alphadog_node/set_rpy', 10)
        self.arm_cmd_pub = self.create_publisher(Float64MultiArray, "/arm_controller/commands", 10)
        self.gripper_pub = self.create_publisher(Float64MultiArray, "/gripper_controller/commands", 10)
        self.get_logger().info("动作发布器已初始化 (发布到 /alphadog_node/* 和 /*_controller/commands)。")

        # ----- ROS 状态订阅者 -----
        self.arm_state_sub = self.create_subscription(
            JointState, "/joint_states", self._arm_state_callback, 10)
        self.get_logger().info(f"订阅机械臂状态 (来自: /joint_states)，期望关节名: {list(self.ros_to_pc_map.keys())}")
        self.dog_state_sub = self.create_subscription(
            BodyStatus, "/alphadog_node/body_status", self._dog_state_callback, 10)
        self.get_logger().info("订阅机器狗状态 (来自: /alphadog_node/body_status)...")
        self.camera_subs = {}
        for cam_name in camera_names:
            topic_name = f"/camera/{cam_name}/image_raw" # TODO: 替换为您实际的摄像头话题
            self.camera_subs[cam_name] = self.create_subscription(
                Image, topic_name, lambda msg, name=cam_name: self._camera_callback(msg, name), 10)
            self.get_logger().info(f"订阅摄像头 (示例: {topic_name})...")
        self.get_logger().info("ROS 桥接节点已初始化。")
        
    def _arm_state_callback(self, msg: JointState):
        with self.lock:
            for i, name in enumerate(msg.name):
                if name in self.ros_to_pc_map:
                    pc_key = self.ros_to_pc_map[name]
                    self._latest_observation[pc_key] = msg.position[i]

    def _dog_state_callback(self, msg: BodyStatus):
        with self.lock:
            self._latest_observation["dog.vx"] = float(msg.vx)
            self._latest_observation["dog.vy"] = float(msg.vy)
            self._latest_observation["dog.wz"] = float(msg.wz)
            self._latest_observation["dog.pitch"] = float(msg.pitch)
            self._latest_observation["dog.roll"] = float(msg.roll)
            self._latest_observation["dog.yaw"] = float(msg.yaw)

    def _camera_callback(self, msg: Image, camera_name: str):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            with self.lock:
                self._latest_frames[camera_name] = cv_image
        except Exception as e:
            self.get_logger().error(f"转换摄像头 '{camera_name}' 图像失败: {e}")

    def send_action(self, action: Dict[str, Any]):
        vel_msg = SetVelocity()
        rpy_msg = SetRpy()
        vel_msg.vx = float(action.get("dog.vx", 0.0))
        vel_msg.vy = float(action.get("dog.vy", 0.0))
        vel_msg.wz = float(action.get("dog.wz", 0.0))
        rpy_msg.pitch = float(action.get("dog.pitch", 0.0))
        rpy_msg.roll = float(action.get("dog.roll", 0.0))
        rpy_msg.yaw = float(action.get("dog.yaw", 0.0))
        self.dog_velocity_pub.publish(vel_msg)
        self.dog_rpy_pub.publish(rpy_msg)
        arm_cmd_msg = Float64MultiArray()
        arm_cmd_msg.data = [float(action.get(key, 0.0)) for key in self.pc_arm_joint_keys]
        self.arm_cmd_pub.publish(arm_cmd_msg)
        gripper_cmd_msg = Float64MultiArray()
        gripper_cmd_msg.data = [float(action.get(self.pc_gripper_key, 0.0))]
        self.gripper_pub.publish(gripper_cmd_msg)

    def get_observation(self) -> Dict[str, Any]:
        with self.lock:
            obs_dict = self._latest_observation.copy()
            frames = {name: img.copy() for name, img in self._latest_frames.items()}
        for cam_key, frame in frames.items():
            ret, buffer = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
            if ret:
                obs_dict[cam_key] = base64.b64encode(buffer).decode("utf-8")
            else:
                obs_dict[cam_key] = ""
        return obs_dict

    def stop_base(self):
        self.get_logger().warn("看门狗超时或程序退出：停止机器狗底盘。")
        self.dog_velocity_pub.publish(SetVelocity())
        self.dog_rpy_pub.publish(SetRpy())


class ZmqHost:
    """ ZMQ 主机，用于管理与 PC 端的通信。"""
    def __init__(self, config: LeDogHostConfig):
        self.zmq_context = zmq.Context()
        self.zmq_cmd_socket = self.zmq_context.socket(zmq.PULL)
        self.zmq_cmd_socket.setsockopt(zmq.CONFLATE, 1)
        self.zmq_cmd_socket.bind(f"tcp://*:{config.port_zmq_cmd}")

        self.zmq_observation_socket = self.zmq_context.socket(zmq.PUSH)
        self.zmq_observation_socket.setsockopt(zmq.CONFLATE, 1)
        self.zmq_observation_socket.bind(f"tcp://*:{config.port_zmq_observations}")

        self.connection_time_s = config.connection_time_s
        self.watchdog_timeout_ms = config.watchdog_timeout_ms
        self.max_loop_freq_hz = config.max_loop_freq_hz
        
        # *** 状态变量，跟踪 PC 连接状态 ***
        self._pc_connected = False
        
        logging.info(f"ZMQ 主机已启动。等待命令端口: {config.port_zmq_cmd} | 观测端口: {config.port_zmq_observations}")
        logging.info("等待 PC 客户端连接...") # 初始提示

    # *** 用于发送观测数据并管理日志 ***
    def send_observation(self, observation_data: str):
        try:
            self.zmq_observation_socket.send_string(observation_data, flags=zmq.NOBLOCK)
            # 如果发送成功，并且之前是断开状态，则打印连接成功信息
            if not self._pc_connected:
                logging.info("PC 客户端已连接，开始发送观测数据。")
                self._pc_connected = True
        except zmq.Again:
            # 如果发送失败 (没有客户端接收)，并且之前是连接状态，则打印断开信息
            if self._pc_connected:
                logging.warning("ZMQ 观测发送失败 (PC 客户端可能已断开)，暂停日志输出...")
                self._pc_connected = False
            # 如果之前已经是断开状态，则不再打印日志

    def disconnect(self):
        self.zmq_observation_socket.close()
        self.zmq_cmd_socket.close()
        self.zmq_context.term()


def main(args=None):
    rclpy.init(args=args)
    
    cfg = LeDogHostConfig()
    
    logging.info("配置 LeDog ROS Bridge 节点")
    robot_node = LeDogRosBridge(cfg.camera_names)

    # 在单独的线程中运行 ROS 节点
    executor = rclpy.executors.MultiThreadedExecutor()
    executor.add_node(robot_node)
    executor_thread = threading.Thread(target=executor.spin, daemon=True)
    executor_thread.start()

    logging.info("启动 ZMQ 主机")
    host = ZmqHost(cfg)

    last_cmd_time = time.time()
    watchdog_active = False
    
    try:
        start = time.perf_counter()
        duration = 0
        while duration < host.connection_time_s and rclpy.ok():
            loop_start_time = time.time()
            
            # 1. (ZMQ) 接收 ZMQ 命令
            try:
                msg = host.zmq_cmd_socket.recv_string(zmq.NOBLOCK)
                data = dict(json.loads(msg))
                # (ROS) 将动作发送到 ROS
                robot_node.send_action(data)
                last_cmd_time = time.time()
                watchdog_active = False
            except zmq.Again:
                pass # 没有收到命令是正常的
            except Exception as e:
                logging.error(f"ZMQ 消息处理失败: {e}")

            # 2. 检查看门狗
            now = time.time()
            if (now - last_cmd_time > host.watchdog_timeout_ms / 1000) and not watchdog_active:
                logging.warning(
                    f"超过 {host.watchdog_timeout_ms}ms 未收到命令。停止底盘。"
                )
                watchdog_active = True
                robot_node.stop_base() # (ROS) 发送停止指令

            # 3. (ROS) 获取观测数据
            last_observation = robot_node.get_observation()

            # 4. (ZMQ) 发送 ZMQ 观测数据
            host.send_observation(json.dumps(last_observation))

            # 5. 循环限频
            elapsed = time.time() - loop_start_time
            time.sleep(max(1 / host.max_loop_freq_hz - elapsed, 0))
            duration = time.perf_counter() - start
            
        print("设定运行时间已到。")

    except KeyboardInterrupt:
        print("收到键盘中断。正在退出...")
    finally:
        print("正在关闭 Lekiwi 主机...")
        robot_node.stop_base()
        host.disconnect()
        executor.shutdown()
        robot_node.destroy_node()
        rclpy.shutdown()
        executor_thread.join()

    logging.info("LeDog OPI 主机已清理并退出。")


if __name__ == "__main__":
    main()